/*=====================================================================*
 *                   Copyright (C) 2011 Paul Mineiro                   *
 * All rights reserved.                                                *
 *                                                                     *
 * Redistribution and use in source and binary forms, with             *
 * or without modification, are permitted provided that the            *
 * following conditions are met:                                       *
 *                                                                     *
 *     * Redistributions of source code must retain the                *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer.                                       *
 *                                                                     *
 *     * Redistributions in binary form must reproduce the             *
 *     above copyright notice, this list of conditions and             *
 *     the following disclaimer in the documentation and/or            *
 *     other materials provided with the distribution.                 *
 *                                                                     *
 *     * Neither the name of Paul Mineiro nor the names                *
 *     of other contributors may be used to endorse or promote         *
 *     products derived from this software without specific            *
 *     prior written permission.                                       *
 *                                                                     *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND              *
 * CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,         *
 * INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES               *
 * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE             *
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER               *
 * OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,                 *
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES            *
 * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE           *
 * GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR                *
 * BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF          *
 * LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT           *
 * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY              *
 * OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE             *
 * POSSIBILITY OF SUCH DAMAGE.                                         *
 *                                                                     *
 * Contact: Paul Mineiro <paul@mineiro.com>                            *
 *=====================================================================*/

#ifndef __NOMINAL_HH__
#define __NOMINAL_HH__

#include <cstring>
#include <vector>

#include "bogackishampine.hh"
#include "fastonebigheader.hh"
#include "model.hh"
#include "normalize.hh"
#include "rating.hh"
#include "pianovec.hh"

namespace playerpiano
{
  class Nominal
    {
      private:
        unsigned int n_workers;
        unsigned int n_labels;
        unsigned int n_items;
        float* alpha;
        float* gamma;
        float prioralpha_sq_stddev;
        unsigned int failcount;

        static inline float
        square (float x)
          {
            return x * x;
          }

      public:
        Nominal (float*       _alpha,
                 float*       _gamma,
                 float        prioralpha_stddev,
                 unsigned int _n_workers,
                 unsigned int _n_labels,
                 unsigned int _n_items)
          : n_workers (_n_workers),
            n_labels (_n_labels),
            n_items (_n_items),
            alpha (_alpha),
            gamma (_gamma),
            prioralpha_sq_stddev (square (prioralpha_stddev)),
            failcount (0)
          {
          }

        Nominal (Model* model)
          : n_workers (1U << strtoul (model->lookup<char*> ("--n_worker_bits").first, 0, 0)),
            n_labels (strtoul (model->lookup<char*> ("--n_labels").first, 0, 0)),
            n_items (strtoul (model->lookup<char*> ("--n_items").first, 0, 0)),
            alpha (model->lookup<float*> ("alpha").first),
            gamma (model->lookup<float*> ("gamma").first),
            prioralpha_sq_stddev (1.0),
            failcount (0)
          {
          }

        void
        apply_deferred_prior (unsigned int  w,
                              uint64_t      s,
                              uint64_t      t,
                              float         eta,
                              float         rho)
          {
            if (t > s)
              {
                float powt = fastpow (t, 1.0f - rho);
                float sigma = prioralpha_sq_stddev; 
                float pows = fastpow (s, 1.0f - rho);
                float decay = fastexp (eta * (pows - powt) / 
                                       (n_items * (1.0f - rho) * sigma));

                float* alphaptr = alpha + w * n_labels * n_labels;

                for (unsigned int j = 0; j < n_labels; ++j)
                  {
                    for (unsigned int k = 0; k < n_labels; ++k)
                      {
                        if (k == j)
                          {
                            /* do nothing */
                          }
                        else
                          {
                            float dalphatmp = 
                              (1.0 - decay) * 
                              (gamma[j * n_labels + k] - 
                               alphaptr[j * n_labels + k]);

                            alphaptr[j * n_labels + k] += dalphatmp;
                            gamma[j * n_labels + k] += 
                              dalphatmp / (1.0 + n_workers);
                          }
                      }
                  }
              }
          }

        class DFunc
          {
            private:
              const Rating&             rating;
              const std::vector<float>& pz;
              const float*              alpha;
              float                     beta;
              unsigned int              n_labels;

            public:
              DFunc (const Rating&             _rating,
                     const std::vector<float>& _pz,
                     const float*              _alpha,
                     float                     _beta,
                     unsigned int              _n_labels)
                  : rating (_rating),
                    pz (_pz),
                    alpha (_alpha),
                    beta (_beta),
                    n_labels (_n_labels)
                {
                }

              template<typename Expr>
              PianoVec 
              operator() (const float&,
                          const Expr& e) const
                {
                  PianoVec rv (n_labels * n_labels, 0);
                  std::vector<float> p (n_labels);
                  const float* alphaptr = 
                    alpha + rating.worker * n_labels * n_labels;

                  for (unsigned int j = 0; j < n_labels; ++j)
                    {
                      for (unsigned int k = 0; k < n_labels; ++k)
                        {
                          if (k == j)
                            {
                              p[k] = beta;
                            }
                          else
                            {
                              p[k] = 
                                -beta * (alphaptr[j * n_labels + k] +
                                         e[j * n_labels + k]);
                            }
                        }

                      safe_exponentiate_normalize (p);

                      for (unsigned int k = 0; k < n_labels; ++k)
                        {
                          if (k == j)
                            {
                              /* do nothing */
                            }
                          else 
                            {
                              rv[j * n_labels + k] += beta * pz[j] * p[k];

                              if (k == rating.label)
                                {
                                  rv[j * n_labels + k] -= beta * pz[j];
                                }
                            }
                        }
                    }

                  return rv;
                }
          };

        void
        update (uint64_t                  t,
                float                     eta,
                float                     rho,
                float                     beta,
                const std::vector<float>& pz,
                const Rating&             rating)
          {
            DFunc dfunc (rating, pz, alpha, beta, n_labels);
            float etat = eta * fastpow (t, -rho);
            PianoVec s (n_labels * n_labels, 0);
            
            std::pair<bool, PianoVec> step = 
              bogacki_shampine_2_3 (dfunc,
                                    0.0f,
                                    s,
                                    etat,
                                    1e-2f * n_labels,
                                    etat / 1000.0,
                                    etat,
                                    1000);

            if (! step.first)
              {
                ++failcount;
              }
            else
              {
                float* alphaptr = alpha + rating.worker * n_labels * n_labels;

                for (unsigned int j = 0; j < n_labels; ++j)
                  {
                    for (unsigned int k = 0; k < n_labels; ++k)
                      {
                        if (k == j)
                          {
                            /* do nothing */
                          }
                        else
                          {
                            alphaptr[j * n_labels + k] += 
                              step.second[j * n_labels + k];
                            gamma[j * n_labels + k] +=
                              step.second[j * n_labels + k] / (1.0 + n_workers);
                          }
                      }
                  }
              }
          }

        float
        logprating (const Rating&   rating,
                    unsigned int    true_label,
                    float           beta) const
          {
            const float* alphaptr = alpha 
                                  + rating.worker * n_labels * n_labels 
                                  + true_label * n_labels;
            unsigned int maxindex = 0;
            float max = (true_label == 0) ? beta : -beta * alphaptr[0];

            for (unsigned int i = 1; i < n_labels; ++i)
              {
                float v = (true_label == i) ? beta : -beta * alphaptr[i];

                if (v > max)
                  {
                    maxindex = i;
                    max = v;
                  }
              }

            float denominator = 0;

            for (unsigned int i = 0; i < n_labels; ++i)
              {
                float v = (true_label == i) ? beta : -beta * alphaptr[i];
                float d = v - max;
                float expd = fastexp (d);
                denominator += expd;
              }

            float v = (true_label == rating.label) ? beta : -beta * alphaptr[rating.label];
            return v - max - fastlog (denominator);
          }

        float
        logprating (const Rating&   rating,
                    unsigned int    true_label,
                    float           beta,
                    float&          dlogbeta,
                    float&          ddlogbeta) const
          {
            const float* alphaptr = alpha 
                                  + rating.worker * n_labels * n_labels 
                                  + true_label * n_labels;
            unsigned int maxindex = 0;
            float max = (true_label == 0) ? beta : -beta * alphaptr[0];

            for (unsigned int i = 1; i < n_labels; ++i)
              {
                float v = (true_label == i) ? beta : -beta * alphaptr[i];

                if (v > max)
                  {
                    maxindex = i;
                    max = v;
                  }
              }

            float denominator = 0;
            float ddenominatordlogbeta = 0;
            float dddenominatorddlogbeta = 0;

            for (unsigned int i = 0; i < n_labels; ++i)
              {
                float v = (true_label == i) ? beta : -beta * alphaptr[i];
                float d = v - max;
                float expd = fastexp (d);
                denominator += expd;
                ddenominatordlogbeta += expd * d;
                dddenominatorddlogbeta += expd * d * (d + 1.0);
              }

            float v = (true_label == rating.label) ? beta : -beta * alphaptr[rating.label];

            dlogbeta = v - max - ddenominatordlogbeta / denominator;
            ddlogbeta = v - max
                        + square (ddenominatordlogbeta / denominator)
                        - dddenominatorddlogbeta / denominator;

            return v - max - fastlog (denominator);
          }

        static void
        ensure (Model* model)
          {
            unsigned int n_workers (1U << strtoul (model->lookup<char*> ("--n_worker_bits").first, 0, 0));
            unsigned int n_labels (strtoul (model->lookup<char*> ("--n_labels").first, 0, 0));
            model->ensure ("alpha",
                           n_workers * square (n_labels) * sizeof (float));
            model->ensure ("gamma", square (n_labels) * sizeof (float));
          }

        unsigned int
        failures () const
          {
            return failcount;
          }
    };
}

#endif // __NOMINAL_HH__
